import os
import caffe
import matplotlib.pyplot as plt
import numpy as np

caffe_root = '/opt/caffe/'

def count_ims():
    path = '/home/xiaomin/wxm/Data/KaggleCCS/test'
    num = 0
    for _ in os.listdir(path):
        if _.endswith('.jpg'):
            num += 1
    print num


def check_caffe_transformer():
    CODE_PATH = '/home/xiaomin/wxm/Code/KaggleCCS'
    model_def = CODE_PATH + '/' + 'prototxt/res50/ResNet-50-deploy.prototxt'
    model_weights = '/home/xiaomin/wxm/Data/KaggleCCS/snapshots/res50/_iter_3000.caffemodel'
    batch_size = 16
    test_image_path = '/home/xiaomin/wxm/Data/KaggleCCS/test'
    submit_file = '/home/xiaomin/wxm/Data/KaggleCCS/submissions/res50/1.csv'
    mu = np.load(caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy')
    mu = mu.mean(1).mean(1)  # average over pixels to obtain the mean (BGR) pixel values

    net = caffe.Net(model_def,  # defines the structure of the model
                    model_weights,  # contains the trained weights
                    caffe.TEST)  # use test mode (e.g., don't perform dropout)

    transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
    transformer.set_transpose('data', (2, 0, 1))  # move image channels to outermost dimension
    transformer.set_mean('data', mu)  # subtract the dataset-mean value in each channel
    # transformer.set_raw_scale('data', 255)  # rescale from [0, 1] to [0, 255]
    # transformer.set_channel_swap('data', (2, 1, 0))  # swap channels from RGB to BGR

    image = caffe.io.load_image('/home/xiaomin/wxm/Data/KaggleCCS/train/Type_1/0.jpg')
    transformed_image = transformer.preprocess('data', image)

    print transformed_image

    # plt.imsave('0.jpg', transformed_image)


if __name__ == '__main__':
    check_caffe_transformer()